package cn.com.cfae.iras.doc.analyze.parse;

import com.hankcs.hanlp.corpus.io.IOUtil;
import com.hankcs.hanlp.mining.word2vec.DocVectorModel;
import com.hankcs.hanlp.mining.word2vec.Word2VecTrainer;
import com.hankcs.hanlp.mining.word2vec.WordVectorModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

public class SimilarAnalyser {

    //    private static final String TRAIN_FILE_NAME =null;// TestUtility.ensureTestData("搜狗文本分类语料库已分词.txt", "http://hanlp.linrunsoft.com/release/corpus/sogou-mini-segmented.zip");
    private static final String MODEL_FILE_NAME = "./data/test/hanlp-wiki-vec-zh.txt";
//    private static final String MODEL_FILE_NAME = "G:\\work\\training\\word-vector_model.txt";
//    private static final String MODEL_FILE_NAME = "data/test/word2vec.txt";

    private static final Logger logger = LoggerFactory.getLogger(SimilarAnalyser.class);

    private static SimilarAnalyser similarAnalyser = new SimilarAnalyser();

    private SimilarAnalyser() {
    }

    public static SimilarAnalyser getInstance() {
        return similarAnalyser;
    }

    private static WordVectorModel wordVectorModel = null;
    private static DocVectorModel docVectorModel;

    static {
        try {
            long t1 = System.currentTimeMillis();
            logger.info("开始加载训练模型......");
            docVectorModel = loadModel();
            long t2 = System.currentTimeMillis();
            logger.info("加载训练模型完成，耗时：{}毫秒。", (t2 - t1));
        } catch (IOException e) {
            logger.error(e.getMessage(), e);
        }
    }


    public double similarity(String textA, String textB) {
        double similarity = -1;
        try {
            similarity = docVectorModel.similarity(textA, textB);
        } catch (Exception e) {
            logger.error(e.getMessage(), e);
        }
        return similarity;
    }


//    static WordVectorModel trainOrLoadModel() throws IOException
//    {
//        if (!IOUtil.isFileExisted(MODEL_FILE_NAME))
//        {
//            if (!IOUtil.isFileExisted(TRAIN_FILE_NAME))
//            {
//                System.err.println("语料不存在，请阅读文档了解语料获取与格式：https://github.com/hankcs/HanLP/wiki/word2vec");
//                System.exit(1);
//            }
//            Word2VecTrainer trainerBuilder = new Word2VecTrainer();
//            return trainerBuilder.train(TRAIN_FILE_NAME, MODEL_FILE_NAME);
//        }
//
//        return loadModel();
//    }

    private static DocVectorModel loadModel() throws IOException {
        logger.info("模型文件目录：{}。", MODEL_FILE_NAME);
        wordVectorModel = new WordVectorModel(MODEL_FILE_NAME);
        docVectorModel = new DocVectorModel(wordVectorModel);
        return docVectorModel;
    }
}
